import os
import argparse
import mxnet.ndarray.random as ndr
from mxnet import profiler
from nn import get_net
import time

def clean(network, root):
    net = get_net(network)
    net.initialize()
    net.hybridize()
    x = ndr.randn(1, 3, net.size, net.size)
    _ = net(x)
    t = time.time()
    for i in range(1000):
        net(x)
    print((time.time() - t))
    if not os.path.exists(root):
        os.makedirs(root)
    net.export(os.path.join(root, network))

def parse_args():
    parser = argparse.ArgumentParser(description='Convert MXNet model to Caffe model')
    parser.add_argument('network', type=str)
    parser.add_argument('--root', type=str, default='models/clean')
    args = parser.parse_args()
    return args

if __name__ == '__main__':
    args = parse_args()
    clean(args.network, args.root)